import numpy

from physics.vector.array import Vector4Array

def four_vectors(nt):
    vecs = {'mup': Vector4Array(nt[:]['mup_cm_px'],
                                nt[:]['mup_cm_py'],
                                nt[:]['mup_cm_pz'],
                                nt[:]['mup_cm_e'],
                                ),
            'mum': Vector4Array(nt[:]['mum_cm_px'],
                                nt[:]['mum_cm_py'],
                                nt[:]['mum_cm_pz'],
                                nt[:]['mum_cm_e'],
                                ),
            'pi01': Vector4Array(nt[:]['pi01_cm_px'],
                                 nt[:]['pi01_cm_py'],
                                 nt[:]['pi01_cm_pz'],
                                 nt[:]['pi01_cm_e'],
                                 ),
            'pi02': Vector4Array(nt[:]['pi02_cm_px'],
                                 nt[:]['pi02_cm_py'],
                                 nt[:]['pi02_cm_pz'],
                                 nt[:]['pi02_cm_e'],
                                 ),
            'initial': Vector4Array(numpy.zeros(len(nt)),
                                    numpy.zeros(len(nt)),
                                    numpy.zeros(len(nt)),
                                    numpy.ones(len(nt)) * 10.876),
            }

    return vecs
    

def m_mumu(nt, vecs):
    return (vecs['mup'] + vecs['mum']).m()

def mm_off_pi0pi0(nt, vecs):
    return (vecs['initial'] - vecs['pi01'] - vecs['pi02']).m()


# 2D fills            
def m_mumu_v_mm_off_pi0pi0(nt, vecs):
    return numpy.array((m_mumu(nt, vecs), mm_off_pi0pi0(nt, vecs))).T

